import torch
from torchvision.datasets import ImageFolder
from utils.dataset import TrainDataset, TestDataset
from torchvision.transforms import transforms
from tqdm import tqdm


def getStat(train_data):
    '''
    Compute mean and variance for training data
    :param train_data: 自定义类Dataset(或ImageFolder即可)
    :return: (mean, std)
    '''
    print('Compute mean and variance for training data.')
    print(len(train_data))
    train_loader = torch.utils.data.DataLoader(
        train_data, batch_size=1, shuffle=False, num_workers=8,
        pin_memory=True)
    mean = torch.zeros(3)
    std = torch.zeros(3)
    for X in tqdm(train_loader):
        for d in range(3):
            mean[d] += X[:, d, :, :].mean()
            std[d] += X[:, d, :, :].std()
    mean.div_(len(train_data))
    std.div_(len(train_data))
    return list(mean.numpy()), list(std.numpy())


if __name__ == '__main__':
    # ([0.6379008, 0.568002, 0.5700531], [0.24534978, 0.25451308, 0.25488153])
    train_dataset = ImageFolder(root=r'../Dataset/trainval', transform=transforms.ToTensor())
    # ([0.6346342, 0.5658491, 0.57049555], [0.24592425, 0.2550813, 0.25544217])
    #
    test_dataset = TestDataset(csv_path='../Dataset/test.csv',
                               file_path='../Dataset/test/',
                               transform=transforms.ToTensor())
    print(getStat(test_dataset))
